#!/usr/bin/env python3
"""
Epic Games DLC解锁器 - 异步事件循环版本
专门处理 /ownership?nsCatalogItemId= 接口
"""

import json
import os
import sys
import re
import asyncio
import subprocess
import winreg
from pathlib import Path
import urllib3
from mitmproxy import options, http
from mitmproxy.tools.dump import DumpMaster

urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

class CertificateManager:
    """证书管理类"""
    
    @staticmethod
    def install_certificate():
        """安装mitmproxy证书"""
        try:
            cert_dir = os.path.expanduser("~/.mitmproxy")
            cert_file = os.path.join(cert_dir, "mitmproxy-ca-cert.pem")
            
            if not os.path.exists(cert_file):
                print("❌ 证书文件不存在，请先运行mitmproxy生成证书")
                return False
            
            print("🔐 正在安装SSL证书...")
            
            result = subprocess.run(
                ['certutil', '-addstore', 'Root', cert_file],
                capture_output=True, 
                text=True, 
                shell=True, 
                timeout=30
            )
            
            if result.returncode == 0:
                print("✅ SSL证书安装成功")
                return True
            else:
                print("❌ 证书安装失败，需要管理员权限")
                return False
                
        except Exception as e:
            print(f"❌ 证书安装错误: {e}")
            return False

class SystemProxy:
    """系统代理管理"""
    
    @staticmethod
    def set_proxy(enable=True):
        """设置系统代理"""
        try:
            key_path = r"Software\Microsoft\Windows\CurrentVersion\Internet Settings"
            with winreg.OpenKey(winreg.HKEY_CURRENT_USER, key_path, 0, winreg.KEY_WRITE) as key:
                
                if enable:
                    try:
                        proxy_enable, _ = winreg.QueryValueEx(key, "ProxyEnable")
                        proxy_server, _ = winreg.QueryValueEx(key, "ProxyServer")
                        SystemProxy.original_settings = {
                            "enabled": proxy_enable,
                            "server": proxy_server
                        }
                    except:
                        SystemProxy.original_settings = {"enabled": 0, "server": ""}
                    
                    winreg.SetValueEx(key, "ProxyEnable", 0, winreg.REG_DWORD, 1)
                    winreg.SetValueEx(key, "ProxyServer", 0, winreg.REG_SZ, "127.0.0.1:8080")
                    winreg.SetValueEx(key, "ProxyOverride", 0, winreg.REG_SZ, "localhost;127.0.0.1")
                    
                    print("✅ 系统代理已启用: 127.0.0.1:8080")
                    return True
                else:
                    if hasattr(SystemProxy, 'original_settings'):
                        if SystemProxy.original_settings["enabled"] == 1:
                            winreg.SetValueEx(key, "ProxyEnable", 0, winreg.REG_DWORD, 1)
                            winreg.SetValueEx(key, "ProxyServer", 0, winreg.REG_SZ, 
                                            SystemProxy.original_settings["server"])
                        else:
                            winreg.SetValueEx(key, "ProxyEnable", 0, winreg.REG_DWORD, 0)
                    
                    print("✅ 系统代理已禁用")
                    return True
                    
        except Exception as e:
            print(f"❌ 代理设置错误: {e}")
            return False

class EpicDLCUnlocker:
    def __init__(self):
        self.modified_count = 0
        
    def response(self, flow: http.HTTPFlow) -> None:
        """处理HTTP响应"""
        url = flow.request.url
        
        if "/epic/ecom/v3/platforms/EPIC/identities/" in url and "/ownership?" in url:
            self._process_ownership_api(flow)
        elif "epicgames.com" in flow.request.host:
            self._process_other_epic_apis(flow)
    
    def _process_ownership_api(self, flow):
        """处理所有权API"""
        print(f"🎯 拦截到Epic所有权接口: {flow.request.path}")
        
        if not flow.response or not flow.response.content:
            return
            
        try:
            data = json.loads(flow.response.content)
            print(f"📦 原始数据: {len(data)} 个DLC项目")
            
            modified = False
            for item in data:
                if "owned" in item and not item["owned"]:
                    item["owned"] = True
                    modified = True
                    print(f"   ✅ 解锁: {item.get('itemId', '未知')}")
            
            if modified:
                flow.response.content = json.dumps(data).encode('utf-8')
                flow.response.headers["content-length"] = str(len(flow.response.content))
                flow.response.headers["x-modified-by"] = "Epic-DLC-Unlocker"
                self.modified_count += 1
                print("🎉 成功解锁DLC项目")
                
                print("📊 修改详情:")
                for item in data:
                    item_id = item.get('itemId', '')[:8] + "..."
                    namespace = item.get('namespace', '')[:8] + "..."
                    print(f"   📁 {namespace} -> {item_id} : {item['owned']}")
            else:
                print("ℹ️ 没有需要修改的项目")
                
        except Exception as e:
            print(f"❌ 处理错误: {e}")
    
    def _process_other_epic_apis(self, flow):
        """处理其他Epic API"""
        if not flow.response or not flow.response.text:
            return
            
        url_lower = flow.request.url.lower()
        if any(k in url_lower for k in ['/entitlements', '/purchases', '/licenses', '/ownership']):
            print(f"🔍 处理Epic API: {flow.request.path}")
            
            try:
                data = json.loads(flow.response.text)
                modified = False
                
                def _modify(obj):
                    nonlocal modified
                    if isinstance(obj, dict):
                        for k, v in obj.items():
                            if k.lower() == "owned" and v is False:
                                obj[k] = True
                                modified = True
                            elif isinstance(v, (dict, list)):
                                _modify(v)
                    elif isinstance(obj, list):
                        for item in obj:
                            _modify(item)
                
                _modify(data)
                
                if modified:
                    flow.response.text = json.dumps(data)
                    self.modified_count += 1
                    print("✅ 修改了Epic API响应")
            except:
                original = flow.response.text
                modified = re.sub(r'"owned"\s*:\s*false', '"owned": true', original, flags=re.IGNORECASE)
                if modified != original:
                    flow.response.text = modified
                    self.modified_count += 1
                    print("✅ 通过文本替换修改了Epic响应")

async def start_proxy():
    """启动代理服务器"""
    opts = options.Options(
        listen_port=8080,
        ssl_insecure=True,
    )
    
    master = DumpMaster(opts)
    master.addons.add(EpicDLCUnlocker())
    
    print("🚀 代理服务器已启动 (127.0.0.1:8080)")
    print("⏹️ 按 Ctrl+C 停止")
    
    await master.run()
    return master

async def run():
    """主运行函数"""
    print("🎮 Epic Games DLC解锁器")
    print("=" * 60)
    print("🎯 专门处理Epic所有权API")
    print("🔧 功能: 自动解锁DLC内容")
    print("=" * 60)
    
    if not CertificateManager.install_certificate():
        print("⚠️ 证书安装可能失败，HTTPS解密可能受限")
    
    if not SystemProxy.set_proxy(enable=True):
        print("⚠️ 代理设置失败")
    
    print("\n💡 使用说明:")
    print("1. 配置Epic客户端使用代理: 127.0.0.1:8080")
    print("2. 启动Epic客户端或游戏")
    print("3. 查看解锁日志")
    print("4. 按 Ctrl+C 停止\n")
    
    try:
        await start_proxy()
    except KeyboardInterrupt:
        print("\n🛑 正在停止代理服务器...")
    finally:
        SystemProxy.set_proxy(enable=False)
        print("✅ 清理完成")

if __name__ == "__main__":
    # 检查管理员权限
    try:
        import ctypes
        if not ctypes.windll.shell32.IsUserAnAdmin():
            print("⚠️ 建议以管理员身份运行")
    except:
        pass
    
    asyncio.run(run())